from contextlib import asynccontextmanager
from fastapi import FastAPI
from config.env import AppConfig
from config.get_db import init_create_table
from config.get_redis import RedisUtil
from config.get_scheduler import SchedulerUtil
from exceptions.handle import handle_exception
from middlewares.handle import handle_middleware
from module_admin.controller.ocr_config_controller import OcrConfigController
from module_admin.controller.llm_controller import llmController
from module_admin.controller.ocr_controller import OcrController
from module_admin.controller.cache_controller import cacheController
from module_admin.controller.captcha_controller import captchaController
from module_admin.controller.common_controller import commonController
from module_admin.controller.config_controller import configController
from module_admin.controller.dept_controller import deptController
from module_admin.controller.dict_controller import dictController
from module_admin.controller.log_controller import logController
from module_admin.controller.login_controller import loginController
from module_admin.controller.job_controller import jobController
from module_admin.controller.menu_controller import menuController
from module_admin.controller.notice_controller import noticeController
from module_admin.controller.online_controller import onlineController
from module_admin.controller.post_controler import postController
from module_admin.controller.role_controller import roleController
from module_admin.controller.server_controller import serverController
from module_admin.controller.user_controller import userController
from module_admin.controller.key_controller import keyController
from module_generator.controller.gen_controller import genController
from sub_applications.handle import handle_sub_applications
from utils.common_util import worship
from utils.log_util import logger
from utils.ocr_llm_util import OcrLlmConfig
from ocr_core.ocr_consumer import OcrConsumer
from ocr_core.llm_consumer import LlmConsumer
from multiprocessing import Process
from multiprocessing import Process
import asyncio
import signal
import sys
from config.mongodb import MongoDB


def run_consumer(
    consumer_type: str,
    consumer_name: str = "Customer",
    ocr_config: dict = {},
):
    """在单独进程中运行消费者（修复事件循环问题）"""

    # 在子进程内创建新的事件循环
    def init_loop():
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        return loop

    async def start_consumer():
        # 创建独立的 Redis 连接
        redis_client = await RedisUtil.create_redis_pool()
        consumer = None
        try:
            if consumer_type == "ocr":
                consumer = OcrConsumer(redis_client, ocr_config)
                logger.info(f"{consumer_name} 启动中...")
                await consumer.start()
            else:
                consumer = LlmConsumer(redis_client)
                logger.info(f"{consumer_name} 启动中...")
                await consumer.start()
        except asyncio.CancelledError:
            # 正常取消，不记录为错误
            logger.info(f"{consumer_name} 被取消")
        except Exception as e:
            logger.error(f"{consumer_name} Consumer 发生错误: {str(e)}")
            raise
        finally:
            if consumer:
                await consumer.stop()
            await redis_client.close()
            logger.info(f"{consumer_name} Consumer 已关闭")

    def handle_signal(signum, frame):
        logger.info(f"收到信号 {signum}，{consumer_name} Consumer 准备关闭...")
        # 停止当前事件循环
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # 取消所有任务
            for task in asyncio.all_tasks(loop):
                task.cancel()
        sys.exit(0)

    # 注册信号处理
    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    try:
        # 创建新的事件循环
        loop = init_loop()
        # 设置当前线程的事件循环
        asyncio.set_event_loop(loop)

        # 在事件循环中运行消费者
        loop.run_until_complete(start_consumer())
    except Exception as e:
        logger.error(f"{consumer_name} Consumer 进程异常退出: {str(e)}")
        sys.exit(1)
    finally:
        # 关闭事件循环
        if loop.is_running():
            loop.stop()
        loop.close()


def generate_consumer_process(consumer_type: str = "llm", num_processes: int = 1):
    processes = []
    for i in range(1, num_processes + 1):
        llm_process = Process(
            target=run_consumer,
            args=(
                consumer_type,
                f"{consumer_type.upper()}-{i}-Consumer",
                app.state.ocr_config,
            ),
            name=f"{consumer_type.upper()}{i}-Consumer",
        )
        llm_process.daemon = True
        llm_process.start()
        processes.append(llm_process)

    return processes


# 生命周期事件
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(f"{AppConfig.app_name} 开始启动")
    worship()
    await init_create_table()
    MongoDB.init()
    app.state.redis = await RedisUtil.create_redis_pool()
    await RedisUtil.init_sys_dict(app.state.redis)
    await RedisUtil.init_sys_config(app.state.redis)
    await RedisUtil.init_api_key_config(app.state.redis)
    await SchedulerUtil.init_system_scheduler()

    # 获取ocr_config和llm_config并缓存
    await OcrLlmConfig.init_llm_config(app.state.redis)
    app.state.ocr_config = await OcrLlmConfig.get_ocr_config()

    # 启动消费者进程
    app.state.consumer_processes = []
    try:
        # 启动1个OCR消费者进程和10个LLM消费者进程
        app.state.consumer_processes += generate_consumer_process("ocr")
        app.state.consumer_processes += generate_consumer_process("llm", 2)

        # 检查进程是否成功启动
        for process in app.state.consumer_processes:
            if not process.is_alive():
                raise RuntimeError(f"{process.name} 启动失败")

        logger.info(f"{AppConfig.app_name} 启动成功")
        yield  # 服务正常运行

    except Exception as e:
        logger.error(f"启动消费者进程时发生错误: {str(e)}")
        # 确保清理所有进程
        for p in app.state.consumer_processes:
            if p.is_alive():
                p.terminate()
                p.join(timeout=5)
        raise
    finally:
        # 关闭时的清理工作
        logger.info("正在关闭消费者进程...")
        for p in app.state.consumer_processes:
            if p.is_alive():
                p.terminate()
                p.join(timeout=5)
                if p.is_alive():
                    p.kill()  # 强制结束
                    p.join()
        await RedisUtil.close_redis_pool(app)
        await SchedulerUtil.close_system_scheduler()
        logger.info("所有服务已关闭")


# 初始化FastAPI对象
app = FastAPI(
    title=AppConfig.app_name,
    description=f"{AppConfig.app_name}接口文档",
    version=AppConfig.app_version,
    lifespan=lifespan,
)

# 挂载子应用
handle_sub_applications(app)
# 加载中间件处理方法
handle_middleware(app)
# 加载全局异常处理方法
handle_exception(app)


# 加载路由列表
controller_list = [
    {"router": loginController, "tags": ["登录模块"]},
    {"router": captchaController, "tags": ["验证码模块"]},
    {"router": userController, "tags": ["系统管理-用户管理"]},
    {"router": roleController, "tags": ["系统管理-角色管理"]},
    {"router": menuController, "tags": ["系统管理-菜单管理"]},
    {"router": deptController, "tags": ["系统管理-部门管理"]},
    {"router": postController, "tags": ["系统管理-岗位管理"]},
    {"router": dictController, "tags": ["系统管理-字典管理"]},
    {"router": configController, "tags": ["系统管理-参数管理"]},
    {"router": noticeController, "tags": ["系统管理-通知公告管理"]},
    {"router": logController, "tags": ["系统管理-日志管理"]},
    {"router": onlineController, "tags": ["系统监控-在线用户"]},
    {"router": jobController, "tags": ["系统监控-定时任务"]},
    {"router": serverController, "tags": ["系统监控-菜单管理"]},
    {"router": cacheController, "tags": ["系统监控-缓存监控"]},
    {"router": commonController, "tags": ["通用模块"]},
    {"router": genController, "tags": ["代码生成模块"]},
    {"router": OcrConfigController, "tags": ["OCR配置"]},
    {"router": OcrController, "tags": ["OCR识别"]},
    {"router": llmController, "tags": ["大模型配置"]},
    {"router": keyController, "tags": ["APIKEY"]},
]

for controller in controller_list:
    app.include_router(router=controller.get("router"), tags=controller.get("tags"))
